import torch
import torch.nn as nn
import numpy as np
import torch.optim as optim
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from Lesson.MLP.lession5_MyNet2 import MyNet2
# python image library python图像处理库
from torchvision import transforms

# 二

img = plt.imread(r"D:\\Photos\\daily\\1655783314442.jpg")
# print(img)
# plt.imshow(img)
# plt.show()
#
# img2 = Image.open(r"D:\\Photos\\daily\\1655783314442.jpg")
# img2.show()
#
# img3 = Image.open(r"D:\\Photos\\daily\\1655783314442.jpg")
# img3 = img3.convert('L')
# img3.show()

# 不推荐这种方式转化为tensor
# imgtsr = torch.tensor(img)
# print(type(imgtsr))
# h w c
print(img.shape)

np.transpose(img, (2, 0, 1))
print(img.shape)

# c h w
imgtr = transforms.ToTensor()(img)
print(type(imgtr))
print(imgtr)



# 一

net2 = MyNet2()

# x = torch.linspace(-7, 7, 200)
# X = x.unsqueeze(1)
# y = torch.cos(x)
# Y = y.unsqueeze(1)

step = 20
x = torch.linspace(1, 20, step)
X = x.unsqueeze(1)
# y = 2*torch.pow(x, 8) + 3*torch.pow(x, 6)
y = x
Y = y.unsqueeze(1)

# 损失函数 MSE
loss_fun = nn.MSELoss()

# 优化器 梯度下降法, lr学习率
optimizer = optim.SGD(net2.parameters(), lr=0.05)
batchSize = 20
c = list(zip(X, Y))
train = DataLoader(c, batch_size=batchSize, shuffle=True)

# 不加激活函数，为什么输出nan？

ls = []
num_epoch = 600
for epoch in range(num_epoch):
    ave = 0.0
    for b_x, b_y in train:
        y_predicted = net2.forward(b_x)
        loss = loss_fun(y_predicted, b_y)
        ave = ave + loss.item()
        optimizer.zero_grad()
        # 损失函数的反向传播
        loss.backward()
        optimizer.step()
    if (epoch+1) % 10 == 0:
        print(f'epoch:{epoch+1}, train_loss:{ave/batchSize:.7f}')
    ls.append(ave/batchSize)

plt.plot(X.numpy(), net2.forward(X).detach().numpy(), 'r')
plt.plot(X.numpy(), Y.numpy(), 'b')
plt.show()

plt.plot(range(num_epoch), ls)
plt.show()

# x = torch.linspace(-7, 7, 200)
# X = x.unsqueeze(1)
# plt.plot(range(200), nn.Sigmoid()(X))
# plt.show()
